sezer91 commited on
Commit
c018080
·
1 Parent(s): 6d91ee0

Add application file

Browse files
Files changed (2) hide show
  1. app.py +31 -43
  2. requirements.txt +1 -1
app.py CHANGED
@@ -1,52 +1,40 @@
1
  import gradio as gr
2
- import torch
3
- from segment_anything import sam_model_registry, SamPredictor
4
  import numpy as np
 
5
  import cv2
 
6
  from PIL import Image
7
 
8
- import os
9
- import urllib.request
10
-
11
- MODEL_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
12
- MODEL_PATH = "sam_vit_b.pth"
13
-
14
- # Eğer model yoksa indir
15
- if not os.path.exists(MODEL_PATH):
16
- print("Model indiriliyor...")
17
- urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
18
- print("Model indirildi.")
19
-
20
- # Model yükle
21
  model_type = "vit_b"
22
-
23
  device = "cuda" if torch.cuda.is_available() else "cpu"
24
- sam = sam_model_registry[model_type](checkpoint=MODEL_PATH)
25
- sam.to(device=device)
26
- predictor = SamPredictor(sam)
27
 
28
- def segment(image, x, y):
29
- image = np.array(image)
30
- predictor.set_image(image)
31
- input_point = np.array([[x, y]])
32
- input_label = np.array([1])
33
- masks, _, _ = predictor.predict(
34
- point_coords=input_point,
35
- point_labels=input_label,
36
- multimask_output=False,
37
- )
38
- mask = masks[0]
39
- masked_image = image.copy()
40
- masked_image[~mask] = 0
41
- return Image.fromarray(masked_image)
42
-
43
- with gr.Blocks() as demo:
44
- with gr.Row():
45
- image_input = gr.Image(type="pil")
46
- x = gr.Number(label="X")
47
- y = gr.Number(label="Y")
48
- btn = gr.Button("Segment")
49
- output = gr.Image()
50
- btn.click(fn=segment, inputs=[image_input, x, y], outputs=output)
51
 
52
- demo.launch(share=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
 
 
2
  import numpy as np
3
+ import torch
4
  import cv2
5
+ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
6
  from PIL import Image
7
 
8
+ # Model yükle
9
+ sam_checkpoint = "sam_vit_b.pth"
 
 
 
 
 
 
 
 
 
 
 
10
  model_type = "vit_b"
 
11
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 
12
 
13
+ sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
14
+ sam.to(device=device)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
 
16
+ mask_generator = SamAutomaticMaskGenerator(sam)
17
+
18
+ def segment_all_objects(image):
19
+ image_np = np.array(image)
20
+ masks = mask_generator.generate(image_np)
21
+
22
+ # Maske üzerine çiz
23
+ overlay = image_np.copy()
24
+ for i, mask in enumerate(masks):
25
+ m = mask["segmentation"]
26
+ color = np.random.randint(0, 255, size=(3,))
27
+ overlay[m] = overlay[m] * 0.3 + color * 0.7
28
+ # Maske üstüne label yaz
29
+ y, x = np.where(m)
30
+ if len(x) > 0 and len(y) > 0:
31
+ cx, cy = int(np.mean(x)), int(np.mean(y))
32
+ cv2.putText(overlay, f"Obj {i+1}", (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
33
+
34
+ return Image.fromarray(overlay.astype(np.uint8))
35
+
36
+ gr.Interface(
37
+ fn=segment_all_objects,
38
+ inputs=gr.Image(type="pil"),
39
+ outputs=gr.Image()
40
+ ).launch()
requirements.txt CHANGED
@@ -1,7 +1,7 @@
1
  gradio
2
  torch
 
3
  opencv-python
4
  numpy
5
  Pillow
6
  git+https://github.com/facebookresearch/segment-anything.git
7
- torchvision
 
1
  gradio
2
  torch
3
+ torchvision
4
  opencv-python
5
  numpy
6
  Pillow
7
  git+https://github.com/facebookresearch/segment-anything.git