EnginDev commited on
Commit
b96e089
·
verified ·
1 Parent(s): a84a1e2

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +443 -0
app.py ADDED
@@ -0,0 +1,443 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import numpy as np
4
+ from PIL import Image
5
+ import cv2
6
+
7
+ print("🚀 Starting SAM2 App v2.0...")
8
+
9
+ device = "cuda" if torch.cuda.is_available() else "cpu"
10
+ print(f"📱 Using device: {device}")
11
+
12
+ model = None
13
+ processor = None
14
+
15
+ def load_model():
16
+ global model, processor
17
+ if model is None:
18
+ print("📦 Loading SAM model...")
19
+ try:
20
+ from transformers import SamModel, SamProcessor
21
+
22
+ model_name = "facebook/sam-vit-large"
23
+
24
+ processor = SamProcessor.from_pretrained(model_name)
25
+ model = SamModel.from_pretrained(model_name)
26
+ model.to(device)
27
+ print(f"✅ Model loaded: {model_name}")
28
+ except Exception as e:
29
+ print(f"❌ Error: {e}, falling back to base model")
30
+ model_name = "facebook/sam-vit-base"
31
+ processor = SamProcessor.from_pretrained(model_name)
32
+ model = SamModel.from_pretrained(model_name)
33
+ model.to(device)
34
+ return model, processor
35
+
36
+ def prepare_image(image, max_size=1024):
37
+ if isinstance(image, np.ndarray):
38
+ image_pil = Image.fromarray(image)
39
+ else:
40
+ image_pil = image
41
+
42
+ if image_pil.mode != 'RGB':
43
+ image_pil = image_pil.convert('RGB')
44
+
45
+ image_np = np.array(image_pil)
46
+ h, w = image_np.shape[:2]
47
+
48
+ if max(h, w) > max_size:
49
+ scale = max_size / max(h, w)
50
+ new_h, new_w = int(h * scale), int(w * scale)
51
+ image_pil = image_pil.resize((new_w, new_h), Image.Resampling.LANCZOS)
52
+ image_np = np.array(image_pil)
53
+
54
+ return image_pil, image_np
55
+
56
+ def refine_mask(mask, kernel_size=5):
57
+ mask_uint8 = (mask > 0).astype(np.uint8) * 255
58
+ kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
59
+ mask_closed = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
60
+ mask_refined = cv2.morphologyEx(mask_closed, cv2.MORPH_OPEN, kernel)
61
+ return mask_refined > 0
62
+
63
+ def merge_nearby_masks(masks, scores, distance_threshold=100):
64
+ """
65
+ Kombiniert Masken die nah beieinander sind
66
+ → Fisch + Flosse werden EINS!
67
+ """
68
+ if len(masks) == 0:
69
+ return np.zeros_like(masks[0]) if len(masks) > 0 else None
70
+
71
+ # Finde Zentren aller Masken
72
+ centers = []
73
+ for mask in masks:
74
+ y_coords, x_coords = np.where(mask)
75
+ if len(x_coords) > 0:
76
+ center_x = int(np.mean(x_coords))
77
+ center_y = int(np.mean(y_coords))
78
+ centers.append((center_x, center_y))
79
+ else:
80
+ centers.append(None)
81
+
82
+ # Gruppiere Masken die nah beieinander sind
83
+ combined = np.zeros_like(masks[0], dtype=bool)
84
+
85
+ for i, mask in enumerate(masks):
86
+ if centers[i] is None:
87
+ continue
88
+
89
+ # Checke ob diese Maske nah an bereits kombinierten Masken ist
90
+ if np.sum(combined) > 0:
91
+ # Distanz zwischen Masken
92
+ combined_coords = np.argwhere(combined)
93
+ mask_coords = np.argwhere(mask)
94
+
95
+ if len(combined_coords) > 0 and len(mask_coords) > 0:
96
+ # Minimale Distanz zwischen den Masken
97
+ min_dist = float('inf')
98
+ for mc in mask_coords[::10]: # Sample alle 10 Punkte
99
+ for cc in combined_coords[::10]:
100
+ dist = np.linalg.norm(mc - cc)
101
+ min_dist = min(min_dist, dist)
102
+
103
+ # Wenn nah genug: kombinieren!
104
+ if min_dist < distance_threshold:
105
+ combined = combined | mask
106
+ print(f" ✅ Merged mask {i} (dist: {min_dist:.1f})")
107
+ else:
108
+ # Erste Maske
109
+ combined = mask
110
+
111
+ return combined
112
+
113
+ def segment_automatic(image, quality="high", merge_parts=True):
114
+ """
115
+ Automatische Segmentierung mit Multi-Point Detection
116
+ merge_parts=True → Kombiniert getrennte Teile wie Flossen!
117
+ """
118
+ if image is None:
119
+ return None, {"error": "Kein Bild hochgeladen"}
120
+
121
+ try:
122
+ print(f"🔄 Starting segmentation (quality: {quality}, merge: {merge_parts})...")
123
+ model, processor = load_model()
124
+
125
+ image_pil, image_np = prepare_image(image)
126
+ h, w = image_np.shape[:2]
127
+
128
+ # MEHRERE Punkte für bessere Abdeckung!
129
+ center_x, center_y = w // 2, h // 2
130
+ points = [
131
+ [center_x, center_y], # Mitte
132
+ [center_x, center_y - h//4], # Oben
133
+ [center_x, center_y + h//4], # Unten
134
+ [center_x - w//4, center_y], # Links
135
+ [center_x + w//4, center_y] # Rechts
136
+ ]
137
+
138
+ all_masks = []
139
+ all_scores = []
140
+
141
+ print(f"🎯 Using {len(points)} detection points...")
142
+
143
+ for idx, point in enumerate(points):
144
+ inputs = processor(
145
+ image_pil,
146
+ input_points=[[point]],
147
+ input_labels=[[1]],
148
+ return_tensors="pt"
149
+ ).to(device)
150
+
151
+ with torch.no_grad():
152
+ outputs = model(**inputs, multimask_output=True)
153
+
154
+ masks = processor.image_processor.post_process_masks(
155
+ outputs.pred_masks.cpu(),
156
+ inputs["original_sizes"].cpu(),
157
+ inputs["reshaped_input_sizes"].cpu()
158
+ )[0]
159
+
160
+ scores = outputs.iou_scores.cpu().numpy()
161
+ if scores.ndim > 1:
162
+ scores = scores.flatten()
163
+
164
+ # Beste Maske von diesem Punkt
165
+ best_idx = np.argmax(scores)
166
+
167
+ if masks.ndim == 4:
168
+ mask = masks[0, best_idx].numpy()
169
+ else:
170
+ mask = masks[best_idx].numpy()
171
+
172
+ # Nur gute Masken
173
+ if scores[best_idx] > 0.6:
174
+ all_masks.append(mask > 0)
175
+ all_scores.append(scores[best_idx])
176
+ print(f" ✅ Point {idx}: Score {scores[best_idx]:.3f}")
177
+
178
+ print(f"✅ Got {len(all_masks)} quality masks")
179
+
180
+ # MERGE: Kombiniere nahe Masken
181
+ if merge_parts and len(all_masks) > 0:
182
+ combined_mask = merge_nearby_masks(all_masks, all_scores, distance_threshold=100)
183
+ else:
184
+ combined_mask = np.zeros((h, w), dtype=bool)
185
+ for mask in all_masks:
186
+ combined_mask = combined_mask | mask
187
+
188
+ # Refinement
189
+ if quality == "high":
190
+ combined_mask = refine_mask(combined_mask, kernel_size=9)
191
+
192
+ # Overlay
193
+ overlay = image_np.copy()
194
+ color = np.array([255, 80, 180]) # Rosa/Pink
195
+
196
+ mask_float = combined_mask.astype(float)
197
+ if quality == "high":
198
+ mask_float = cv2.GaussianBlur(mask_float, (7, 7), 0)
199
+
200
+ for c in range(3):
201
+ overlay[:, :, c] = (
202
+ overlay[:, :, c] * (1 - mask_float * 0.65) +
203
+ color[c] * mask_float * 0.65
204
+ )
205
+
206
+ # Kontur
207
+ contours, _ = cv2.findContours(
208
+ combined_mask.astype(np.uint8),
209
+ cv2.RETR_EXTERNAL,
210
+ cv2.CHAIN_APPROX_SIMPLE
211
+ )
212
+ cv2.drawContours(overlay, contours, -1, (255, 255, 0), 3)
213
+
214
+ metadata = {
215
+ "success": True,
216
+ "mode": "automatic_plus" if merge_parts else "automatic",
217
+ "quality": quality,
218
+ "detection_points": len(points),
219
+ "masks_found": len(all_masks),
220
+ "avg_score": float(np.mean(all_scores)) if all_scores else 0,
221
+ "all_scores": [float(s) for s in all_scores],
222
+ "image_size": [w, h],
223
+ "mask_area": int(np.sum(combined_mask)),
224
+ "mask_percentage": float(np.sum(combined_mask) / (h * w) * 100),
225
+ "num_contours": len(contours),
226
+ "device": device
227
+ }
228
+
229
+ print(f"✅ Segmentation complete! {len(contours)} contours")
230
+ return Image.fromarray(overlay.astype(np.uint8)), metadata
231
+
232
+ except Exception as e:
233
+ import traceback
234
+ print(f"❌ ERROR:\n{traceback.format_exc()}")
235
+ return image, {"error": str(e)}
236
+
237
+ def segment_multi_dense(image, density="medium"):
238
+ """Multi-Object mit dichterem Grid"""
239
+ if image is None:
240
+ return None, {"error": "Kein Bild"}
241
+
242
+ try:
243
+ model, processor = load_model()
244
+ image_pil, image_np = prepare_image(image)
245
+ h, w = image_np.shape[:2]
246
+
247
+ if density == "high":
248
+ grid_size = 5
249
+ elif density == "medium":
250
+ grid_size = 4
251
+ else:
252
+ grid_size = 3
253
+
254
+ points = []
255
+ for i in range(1, grid_size + 1):
256
+ for j in range(1, grid_size + 1):
257
+ x = int(w * i / (grid_size + 1))
258
+ y = int(h * j / (grid_size + 1))
259
+ points.append([x, y])
260
+
261
+ print(f"📍 Using {len(points)} points for segmentation")
262
+
263
+ all_masks = []
264
+ all_scores = []
265
+
266
+ for idx, point in enumerate(points):
267
+ inputs = processor(
268
+ image_pil,
269
+ input_points=[[point]],
270
+ input_labels=[[1]],
271
+ return_tensors="pt"
272
+ ).to(device)
273
+
274
+ with torch.no_grad():
275
+ outputs = model(**inputs, multimask_output=True)
276
+
277
+ masks = processor.image_processor.post_process_masks(
278
+ outputs.pred_masks.cpu(),
279
+ inputs["original_sizes"].cpu(),
280
+ inputs["reshaped_input_sizes"].cpu()
281
+ )[0]
282
+
283
+ scores = outputs.iou_scores.cpu().numpy().flatten()
284
+ best_idx = np.argmax(scores)
285
+
286
+ if masks.ndim == 4:
287
+ mask = masks[0, best_idx].numpy()
288
+ else:
289
+ mask = masks[best_idx].numpy()
290
+
291
+ if scores[best_idx] > 0.7:
292
+ all_masks.append(refine_mask(mask))
293
+ all_scores.append(scores[best_idx])
294
+
295
+ print(f"✅ Got {len(all_masks)} quality masks")
296
+
297
+ overlay = image_np.copy()
298
+
299
+ colors = []
300
+ for i in range(len(all_masks)):
301
+ hue = int(180 * i / max(len(all_masks), 1))
302
+ color_hsv = np.uint8([[[hue, 255, 200]]])
303
+ color_rgb = cv2.cvtColor(color_hsv, cv2.COLOR_HSV2RGB)[0][0]
304
+ colors.append(color_rgb)
305
+
306
+ for mask, color, score in zip(all_masks, colors, all_scores):
307
+ alpha = 0.4 + (score - 0.7) * 0.2
308
+ overlay[mask] = (
309
+ overlay[mask] * (1 - alpha) +
310
+ np.array(color) * alpha
311
+ ).astype(np.uint8)
312
+
313
+ contours, _ = cv2.findContours(
314
+ mask.astype(np.uint8),
315
+ cv2.RETR_EXTERNAL,
316
+ cv2.CHAIN_APPROX_SIMPLE
317
+ )
318
+ cv2.drawContours(overlay, contours, -1, color.tolist(), 2)
319
+
320
+ metadata = {
321
+ "success": True,
322
+ "mode": "multi_object_dense",
323
+ "density": density,
324
+ "grid_size": f"{grid_size}x{grid_size}",
325
+ "total_points": len(points),
326
+ "quality_masks": len(all_masks),
327
+ "avg_score": float(np.mean(all_scores)) if all_scores else 0,
328
+ "scores": [float(s) for s in all_scores]
329
+ }
330
+
331
+ return Image.fromarray(overlay), metadata
332
+
333
+ except Exception as e:
334
+ import traceback
335
+ print(f"❌ ERROR:\n{traceback.format_exc()}")
336
+ return image, {"error": str(e)}
337
+
338
+ # Gradio Interface
339
+ demo = gr.Blocks(title="SAM2 Boostly", theme=gr.themes.Soft())
340
+
341
+ with demo:
342
+ gr.Markdown("# 🎨 SAM2 Segmentierung - Boostly Edition")
343
+ gr.Markdown("### Zero-Shot Object Segmentation powered by Meta SAM")
344
+
345
+ with gr.Tab("🤖 Automatisch PLUS"):
346
+ gr.Markdown("**Intelligent Multi-Point Detection** - Erkennt und kombiniert alle Objektteile!")
347
+ gr.Markdown("✨ Perfekt für: Fische mit Flossen, Menschen mit Armen/Beinen, komplexe Objekte")
348
+
349
+ with gr.Row():
350
+ with gr.Column():
351
+ input_auto = gr.Image(type="pil", label="📸 Bild hochladen")
352
+
353
+ quality_radio = gr.Radio(
354
+ choices=["high", "fast"],
355
+ value="high",
356
+ label="Qualität"
357
+ )
358
+
359
+ merge_checkbox = gr.Checkbox(
360
+ value=True,
361
+ label="✅ Teile zusammenfügen (Empfohlen!)"
362
+ )
363
+
364
+ btn_auto = gr.Button("🚀 Segmentieren", variant="primary", size="lg")
365
+
366
+ gr.Markdown("""
367
+ **So funktioniert's:**
368
+ - 🎯 5 Erkennungspunkte (Mitte + 4 Richtungen)
369
+ - 🔗 Intelligentes Merging naher Bereiche
370
+ - 🎨 Fisch + Flosse = EIN zusammenhängendes Objekt!
371
+ """)
372
+
373
+ with gr.Column():
374
+ output_auto = gr.Image(label="✨ Ergebnis")
375
+ json_auto = gr.JSON(label="📊 Details")
376
+
377
+ btn_auto.click(
378
+ fn=segment_automatic,
379
+ inputs=[input_auto, quality_radio, merge_checkbox],
380
+ outputs=[output_auto, json_auto]
381
+ )
382
+
383
+ with gr.Tab("🎯 Multi-Region"):
384
+ gr.Markdown("**Grid-basierte Segmentierung** - Für mehrere separate Objekte")
385
+
386
+ with gr.Row():
387
+ with gr.Column():
388
+ input_multi = gr.Image(type="pil", label="📸 Bild hochladen")
389
+ density_radio = gr.Radio(
390
+ choices=["high", "medium", "low"],
391
+ value="medium",
392
+ label="Punkt-Dichte"
393
+ )
394
+ btn_multi = gr.Button("🎯 Alle Bereiche", variant="primary", size="lg")
395
+
396
+ gr.Markdown("""
397
+ **Grid-Größen:**
398
+ - 🔥 High: 5x5 = 25 Punkte
399
+ - ⚡ Medium: 4x4 = 16 Punkte
400
+ - 💨 Low: 3x3 = 9 Punkte
401
+ """)
402
+
403
+ with gr.Column():
404
+ output_multi = gr.Image(label="✨ Ergebnis")
405
+ json_multi = gr.JSON(label="📊 Details")
406
+
407
+ btn_multi.click(
408
+ fn=segment_multi_dense,
409
+ inputs=[input_multi, density_radio],
410
+ outputs=[output_multi, json_multi]
411
+ )
412
+
413
+ with gr.Tab("📡 API Integration"):
414
+ gr.Markdown("### API Endpoint für Lovable/externe Apps")
415
+ gr.Code(
416
+ "https://EnginDev-Boostly.hf.space/api/predict",
417
+ label="API URL"
418
+ )
419
+
420
+ gr.Markdown("### JavaScript Beispiel")
421
+ gr.Code('''
422
+ const response = await fetch('https://EnginDev-Boostly.hf.space/api/predict', {
423
+ method: 'POST',
424
+ headers: {'Content-Type': 'application/json'},
425
+ body: JSON.stringify({
426
+ data: [imageBase64, "high", true], // [image, quality, merge]
427
+ fn_index: 0 // 0 = Automatisch PLUS, 1 = Multi-Region
428
+ })
429
+ });
430
+ const result = await response.json();
431
+ const segmentedImage = result.data[0]; // Base64 Image
432
+ const metadata = result.data[1]; // JSON Details
433
+ ''', language="javascript")
434
+
435
+ gr.Markdown("### Parameter")
436
+ gr.Markdown("- **fn_index**: 0 = Automatisch PLUS, 1 = Multi-Region")
437
+ gr.Markdown("- **quality**: 'high' (präzise) oder 'fast' (schnell)")
438
+ gr.Markdown("- **merge**: true (Teile kombinieren) oder false")
439
+ gr.Markdown("- **density**: 'high', 'medium', 'low' (nur Multi-Region)")
440
+
441
+ if __name__ == "__main__":
442
+ print("🌐 Launching Boostly SAM2...")
443
+ demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)