CallmeKaito commited on
Commit
88fd830
1 Parent(s): 10d1571

Delete SAM.py

Browse files
Files changed (1) hide show
  1. SAM.py +0 -205
SAM.py DELETED
@@ -1,205 +0,0 @@
1
- #!/usr/bin/env python
2
- # coding: utf-8
3
-
4
- # # Utility functions
5
-
6
- # In[ ]:
7
-
8
-
9
- import numpy as np
10
- import matplotlib.pyplot as plt
11
-
12
- def show_mask(mask, ax, random_color=False):
13
- if random_color:
14
- color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
15
- else:
16
- color = np.array([30/255, 144/255, 255/255, 0.6])
17
- h, w = mask.shape[-2:]
18
- mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
19
- ax.imshow(mask_image)
20
-
21
-
22
- def show_box(box, ax):
23
- x0, y0 = box[0], box[1]
24
- w, h = box[2] - box[0], box[3] - box[1]
25
- ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))
26
-
27
- def show_boxes_on_image(raw_image, boxes):
28
- plt.figure(figsize=(10,10))
29
- plt.imshow(raw_image)
30
- for box in boxes:
31
- show_box(box, plt.gca())
32
- plt.axis('on')
33
- plt.show()
34
-
35
- def show_points_on_image(raw_image, input_points, input_labels=None):
36
- plt.figure(figsize=(10,10))
37
- plt.imshow(raw_image)
38
- input_points = np.array(input_points)
39
- if input_labels is None:
40
- labels = np.ones_like(input_points[:, 0])
41
- else:
42
- labels = np.array(input_labels)
43
- show_points(input_points, labels, plt.gca())
44
- plt.axis('on')
45
- plt.show()
46
-
47
- def show_points_and_boxes_on_image(raw_image, boxes, input_points, input_labels=None):
48
- plt.figure(figsize=(10,10))
49
- plt.imshow(raw_image)
50
- input_points = np.array(input_points)
51
- if input_labels is None:
52
- labels = np.ones_like(input_points[:, 0])
53
- else:
54
- labels = np.array(input_labels)
55
- show_points(input_points, labels, plt.gca())
56
- for box in boxes:
57
- show_box(box, plt.gca())
58
- plt.axis('on')
59
- plt.show()
60
-
61
-
62
- def show_points_and_boxes_on_image(raw_image, boxes, input_points, input_labels=None):
63
- plt.figure(figsize=(10,10))
64
- plt.imshow(raw_image)
65
- input_points = np.array(input_points)
66
- if input_labels is None:
67
- labels = np.ones_like(input_points[:, 0])
68
- else:
69
- labels = np.array(input_labels)
70
- show_points(input_points, labels, plt.gca())
71
- for box in boxes:
72
- show_box(box, plt.gca())
73
- plt.axis('on')
74
- plt.show()
75
-
76
-
77
- def show_points(coords, labels, ax, marker_size=375):
78
- pos_points = coords[labels==1]
79
- neg_points = coords[labels==0]
80
- ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
81
- ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
82
-
83
-
84
- def show_masks_on_image(raw_image, masks, scores):
85
- if len(masks.shape) == 4:
86
- masks = masks.squeeze()
87
- if scores.shape[0] == 1:
88
- scores = scores.squeeze()
89
-
90
- nb_predictions = scores.shape[-1]
91
- fig, axes = plt.subplots(1, nb_predictions, figsize=(15, 15))
92
-
93
- for i, (mask, score) in enumerate(zip(masks, scores)):
94
- mask = mask.cpu().detach()
95
- axes[i].imshow(np.array(raw_image))
96
- show_mask(mask, axes[i])
97
- axes[i].title.set_text(f"Mask {i+1}, Score: {score.item():.3f}")
98
- axes[i].axis("off")
99
- plt.show()
100
-
101
-
102
- # # Model loading
103
-
104
- # In[ ]:
105
-
106
-
107
- import torch
108
- from transformers import SamModel, SamProcessor
109
-
110
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
111
- model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
112
- processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
113
-
114
-
115
- # In[ ]:
116
-
117
-
118
- from PIL import Image
119
- import requests
120
-
121
- img_url = "thuya.jpeg"
122
- raw_image = Image.open(img_url)
123
-
124
- plt.imshow(raw_image)
125
-
126
-
127
- # ## Step 1: Retrieve the image embeddings
128
-
129
- # In[ ]:
130
-
131
-
132
- inputs = processor(raw_image, return_tensors="pt").to(device)
133
- image_embeddings = model.get_image_embeddings(inputs["pixel_values"])
134
-
135
-
136
- # In[ ]:
137
-
138
-
139
- input_points = [[[200, 300]]]
140
- show_points_on_image(raw_image, input_points[0])
141
-
142
-
143
- # In[ ]:
144
-
145
-
146
- inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(device)
147
- # pop the pixel_values as they are not neded
148
- inputs.pop("pixel_values", None)
149
- inputs.update({"image_embeddings": image_embeddings})
150
-
151
- with torch.no_grad():
152
- outputs = model(**inputs)
153
-
154
- masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
155
- scores = outputs.iou_scores
156
-
157
-
158
- # In[ ]:
159
-
160
-
161
- show_masks_on_image(raw_image, masks[0], scores)
162
-
163
-
164
- # ## Export the masked images
165
-
166
- # In[92]:
167
-
168
-
169
- import cv2
170
-
171
- if len(masks[0].shape) == 4:
172
- masks[0] = masks[0].squeeze()
173
- if scores.shape[0] == 1:
174
- scores = scores.squeeze()
175
-
176
- nb_predictions = scores.shape[-1]
177
- fig, axes = plt.subplots(1, nb_predictions, figsize=(15, 15))
178
- for i, (mask, score) in enumerate(zip(masks[0], scores)):
179
- mask = mask.cpu().detach()
180
- axes[i].imshow(np.array(raw_image))
181
- # show_mask(mask, axes[i])
182
-
183
- mask_image = (mask.numpy() * 255).astype(np.uint8) # Convert to uint8 format
184
- cv2.imwrite('mask.png', mask_image)
185
-
186
- image = cv2.imread('thuya.jpeg')
187
-
188
- color_mask = np.zeros_like(image)
189
- color_mask[mask > 0.5] = [30, 144, 255] # Choose any color you like
190
- masked_image = cv2.addWeighted(image, 0.6, color_mask, 0.4, 0)
191
-
192
- color = np.array([30/255, 144/255, 255/255])
193
- #mask_image = * color.reshape(1, 1, -1)
194
-
195
- new_image = -image* np.tile(mask_image[...,None], 3)
196
-
197
- cv2.imwrite('masked_image2.png', cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR))
198
-
199
-
200
-
201
- # In[85]:
202
-
203
-
204
- .shape
205
-