IZERE HIRWA Roger commited on
Commit
cbb1938
·
1 Parent(s): 7d4aa82
Files changed (1) hide show
  1. app.py +7 -4
app.py CHANGED
@@ -39,13 +39,16 @@ app = Flask(__name__)
39
  CORS(app)
40
 
41
  def segment(image_pil: Image.Image, prompt: str):
42
- # Convert PIL image to numpy array
43
- image_np = np.array(image_pil)
 
 
 
44
 
45
  # Run GroundingDINO to get boxes for the prompt
46
  boxes, _, _ = predict(
47
  model=grounder,
48
- image=torch.tensor(image_np).to(device), # Ensure image is on CPU
49
  caption=prompt,
50
  box_threshold=0.3,
51
  text_threshold=0.25,
@@ -56,7 +59,7 @@ def segment(image_pil: Image.Image, prompt: str):
56
 
57
  # 2) Largest box → mask via SAM
58
  box = boxes[np.argmax((boxes[:,2]-boxes[:,0])*(boxes[:,3]-boxes[:,1]))]
59
- predictor.set_image(image_np)
60
  masks, _, _ = predictor.predict(box=box)
61
  mask = masks[0] # boolean HxW
62
 
 
39
  CORS(app)
40
 
41
  def segment(image_pil: Image.Image, prompt: str):
42
+ # Convert PIL image to numpy array and normalize
43
+ image_np = np.array(image_pil).astype(np.float32) / 255.0 # Normalize to [0, 1]
44
+
45
+ # Convert numpy array to torch tensor
46
+ image_tensor = torch.tensor(image_np).permute(2, 0, 1).unsqueeze(0).to(device) # Convert to CHW format
47
 
48
  # Run GroundingDINO to get boxes for the prompt
49
  boxes, _, _ = predict(
50
  model=grounder,
51
+ image=image_tensor, # Pass normalized tensor
52
  caption=prompt,
53
  box_threshold=0.3,
54
  text_threshold=0.25,
 
59
 
60
  # 2) Largest box → mask via SAM
61
  box = boxes[np.argmax((boxes[:,2]-boxes[:,0])*(boxes[:,3]-boxes[:,1]))]
62
+ predictor.set_image(np.array(image_pil))
63
  masks, _, _ = predictor.predict(box=box)
64
  mask = masks[0] # boolean HxW
65