ArslanRobo commited on
Commit
ac614bf
·
verified ·
1 Parent(s): 5ba7e2c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +130 -0
app.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """We can use Gradio to build the UI and then make it compatible for the Hugging face."""
2
+ import gradio as gr
3
+ import cv2
4
+ import numpy as np
5
+ import imutils
6
+ from PIL import Image
7
+
8
+ cv2.ocl.setUseOpenCL(False)
9
+
10
+ # Sharpening function
11
+ def image_sharpening(image):
12
+ kernel_sharpening = np.array([[-1, -1, -1],
13
+ [-1, 9, -1],
14
+ [-1, -1, -1]])
15
+ sharpened = cv2.filter2D(image, -1, kernel_sharpening)
16
+ return sharpened
17
+
18
+ # Remove black borders function
19
+ def remove_black_region(result):
20
+ gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
21
+ thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)[1]
22
+ cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
23
+ cnts = imutils.grab_contours(cnts)
24
+ c = max(cnts, key=cv2.contourArea)
25
+ (x, y, w, h) = cv2.boundingRect(c)
26
+ crop = result[y:y + h, x:x + w]
27
+ return crop
28
+
29
+ # Key point detection and descriptor function
30
+ def detectAndDescribe(image, method='orb'):
31
+ if method == 'sift':
32
+ descriptor = cv2.SIFT_create()
33
+ elif method == 'brisk':
34
+ descriptor = cv2.BRISK_create()
35
+ elif method == 'orb':
36
+ descriptor = cv2.ORB_create()
37
+ (kps, features) = descriptor.detectAndCompute(image, None)
38
+ return kps, features
39
+
40
+ # Matcher creation
41
+ def createMatcher(method, crossCheck):
42
+ if method in ['sift', 'surf']:
43
+ bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=crossCheck)
44
+ else:
45
+ bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=crossCheck)
46
+ return bf
47
+
48
+ # Matching key points
49
+ def matchKeyPointsKNN(featuresA, featuresB, ratio, method):
50
+ bf = createMatcher(method, crossCheck=False)
51
+ rawMatches = bf.knnMatch(featuresA, featuresB, 2)
52
+ matches = []
53
+ for m, n in rawMatches:
54
+ if m.distance < n.distance * ratio:
55
+ matches.append(m)
56
+ return matches
57
+
58
+ # Homography calculation
59
+ def getHomography(kpsA, kpsB, featuresA, featuresB, matches, reprojThresh=4.0):
60
+ kpsA = np.float32([kp.pt for kp in kpsA])
61
+ kpsB = np.float32([kp.pt for kp in kpsB])
62
+ if len(matches) > 4:
63
+ ptsA = np.float32([kpsA[m.queryIdx] for m in matches])
64
+ ptsB = np.float32([kpsB[m.trainIdx] for m in matches])
65
+ (H, status) = cv2.findHomography(ptsA, ptsB, cv2.RANSAC, reprojThresh)
66
+ return matches, H, status
67
+ else:
68
+ return None
69
+
70
+ # Stitching function for two images
71
+ def stitch_two_images(queryImg, trainImg, feature_extractor):
72
+ queryImg_gray = cv2.cvtColor(queryImg, cv2.COLOR_BGR2GRAY)
73
+ trainImg_gray = cv2.cvtColor(trainImg, cv2.COLOR_BGR2GRAY)
74
+ kpsA, featuresA = detectAndDescribe(trainImg_gray, method=feature_extractor)
75
+ kpsB, featuresB = detectAndDescribe(queryImg_gray, method=feature_extractor)
76
+ matches = matchKeyPointsKNN(featuresA, featuresB, ratio=0.75, method=feature_extractor)
77
+ M = getHomography(kpsA, kpsB, featuresA, featuresB, matches, reprojThresh=5)
78
+ if M is None:
79
+ return None
80
+ (matches, H, status) = M
81
+ width = trainImg.shape[1] + queryImg.shape[1]
82
+ height = trainImg.shape[0] + queryImg.shape[0]
83
+ result = cv2.warpPerspective(trainImg, H, (width, height))
84
+ result[0:queryImg.shape[0], 0:queryImg.shape[1]] = queryImg
85
+ crop_image = remove_black_region(result)
86
+ return crop_image
87
+
88
+ # Calculate target brightness
89
+ def calculate_target_brightness(images):
90
+ brightness_values = [np.mean(image.astype(np.float32)) for image in images]
91
+ return np.mean(brightness_values)
92
+
93
+ # Brightness adjustment
94
+ def global_brightness_adjustment(images, target_brightness):
95
+ adjusted_images = []
96
+ for image in images:
97
+ image_float = image.astype(np.float32)
98
+ avg_brightness = np.mean(image_float)
99
+ brightness_shift = target_brightness - avg_brightness
100
+ adjusted_image = image_float + brightness_shift
101
+ adjusted_image = np.clip(adjusted_image, 0, 255).astype(np.uint8)
102
+ adjusted_images.append(adjusted_image)
103
+ return adjusted_images
104
+
105
+ # Main Stitching function
106
+ def stitch_images(uploaded_files, feature_extractor):
107
+ images = [cv2.cvtColor(np.array(Image.open(file)), cv2.COLOR_RGB2BGR) for file in uploaded_files]
108
+ if len(images) == 0:
109
+ return None
110
+ # feature_extractor = 'orb'
111
+ target_brightness = calculate_target_brightness(images)
112
+ adjusted_images = global_brightness_adjustment(images, target_brightness)
113
+ stitched_image = adjusted_images[0]
114
+ for i in range(1, len(adjusted_images)):
115
+ queryImg = stitched_image
116
+ trainImg = adjusted_images[i]
117
+ stitched_image = stitch_two_images(queryImg, trainImg, feature_extractor)
118
+ return cv2.cvtColor(stitched_image, cv2.COLOR_BGR2RGB)
119
+
120
+ # Gradio interface with feature extractor selector
121
+ with gr.Blocks() as demo:
122
+ gr.Markdown("## Image Stitching App with Feature Extractor Selection")
123
+ image_input = gr.Files(label="Upload Images", type="filepath")
124
+ extractor_input = gr.Dropdown(choices=["orb", "sift", "brisk"], label="Feature Extractor", value="orb")
125
+ image_output = gr.Image(type="numpy", label="Stitched Image")
126
+ process_button = gr.Button("Process Image")
127
+ process_button.click(stitch_images, inputs=[image_input, extractor_input], outputs=image_output)
128
+
129
+ # Launch the Gradio app
130
+ demo.launch()