ppicazo commited on
Commit
b920e42
·
verified ·
1 Parent(s): 5c00058

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -10
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
- from PIL import Image, ImageDraw, ImageFont
4
 
5
  # Load object detection pipeline
6
  model_pipeline = pipeline(
@@ -14,30 +14,41 @@ def predict(image):
14
  height = int(image.height * ratio)
15
  image = image.resize((width, height))
16
 
17
- detections = model_pipeline(image)
18
 
19
- # Draw boxes
20
  draw = ImageDraw.Draw(image)
 
 
21
  for det in detections:
22
  box = det["box"]
23
- label = f'{det["label"]} ({det["score"]:.2f})'
24
- print(f"Drawing box: {box} Label: {label}") # Debug
 
 
 
 
 
 
 
 
25
 
26
- # Draw rectangle
27
  draw.rectangle(
28
  [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
29
  outline="red",
30
  width=3
31
  )
32
- # Optional: draw label
33
- draw.text((box["xmin"] + 2, box["ymin"] - 10), label, fill="red")
34
 
35
- return image
36
 
 
37
  gr.Interface(
38
  fn=predict,
39
  inputs=gr.Image(type="pil", label="Upload Astrophotography Image"),
40
- outputs=gr.Image(type="pil", label="Detected Objects"),
 
 
 
41
  title="Astrophotography Object Detector",
42
  allow_flagging="manual",
43
  ).launch()
 
1
  import gradio as gr
2
  from transformers import pipeline
3
+ from PIL import Image, ImageDraw
4
 
5
  # Load object detection pipeline
6
  model_pipeline = pipeline(
 
14
  height = int(image.height * ratio)
15
  image = image.resize((width, height))
16
 
17
+ detections = model_pipeline(image, threshold=0.1)
18
 
 
19
  draw = ImageDraw.Draw(image)
20
+ table_rows = []
21
+
22
  for det in detections:
23
  box = det["box"]
24
+ label = det["label"]
25
+ score = round(det["score"], 4)
26
+ table_rows.append({
27
+ "Class": label,
28
+ "Confidence": f"{score:.2%}",
29
+ "Xmin": int(box["xmin"]),
30
+ "Ymin": int(box["ymin"]),
31
+ "Xmax": int(box["xmax"]),
32
+ "Ymax": int(box["ymax"]),
33
+ })
34
 
 
35
  draw.rectangle(
36
  [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
37
  outline="red",
38
  width=3
39
  )
40
+ draw.text((box["xmin"] + 4, box["ymin"] - 12), f"{label} ({score:.2f})", fill="red")
 
41
 
42
+ return image, table_rows
43
 
44
+ # Gradio Interface
45
  gr.Interface(
46
  fn=predict,
47
  inputs=gr.Image(type="pil", label="Upload Astrophotography Image"),
48
+ outputs=[
49
+ gr.Image(type="pil", label="Detected Objects"),
50
+ gr.Dataframe(headers=["Class", "Confidence", "Xmin", "Ymin", "Xmax", "Ymax"], label="Detections")
51
+ ],
52
  title="Astrophotography Object Detector",
53
  allow_flagging="manual",
54
  ).launch()