Update app.py

app.py

@@ -1,33 +1,16 @@
 import gradio as gr
-from transformers import AutoModelForImageSegmentation, AutoTokenizer
 import torch
-from PIL import Image
+from ultralytics import YOLO
 
-# Load the model
-model = AutoModelForImageSegmentation.from_pretrained("/absolute/path/to/gelan-c-seg.pt")
+# Load YOLO model
+model = YOLO("yolov9c-seg.pt")
 
-# Load the tokenizer (if needed)
-tokenizer = AutoTokenizer.from_pretrained("path/to/tokenizer")
+# Define function to perform prediction
+def predict_image(image):
+    result = model.predict(image, show=True, save=True)
+    return result
 
-# Define the prediction function
-def predict_segmentation(image):
-    # Convert image to PyTorch tensor
-    image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0).float()
-    # Perform inference
-    output = model(image)
-    # Process the output as needed (e.g., post-processing for segmentation masks)
-    
-    segmentation_mask = output.logits.argmax(dim=1).squeeze().detach().numpy()
-    return segmentation_mask
+# Create Gradio interface
+image_input = gr.inputs.Image()
+gr.Interface(predict_image, image_input, "image").launch()
 
-# Create a Gradio interface
-inputs = gr.inputs.Image(shape=(224, 224))
-outputs = gr.outputs.Image(type="numpy", label="Segmentation Mask")
-title = "Image Segmentation Demo"
-description = "Upload an image and get the segmentation mask."
-examples = [["example.jpg"]]  # Add example images here if needed
-interface = gr.Interface(fn=predict_segmentation, inputs=inputs, outputs=outputs, title=title, description=description, examples=examples)
-
-# Run the interface
-if __name__ == "__main__":
-    interface.launch()