Update app.py

app.py

@@ -4,11 +4,21 @@ import numpy as np
 import cv2
 from PIL import Image
 import random
+from transformers import pipeline
 
-# Load YOLOv8 segmentation model (nano version for speed)
+# ---------------------------
+# Load Models
+# ---------------------------
+
+# Text model (tiny LLM)
+text_gen = pipeline("text-generation", model="tiny-random-gpt2")
+
+# YOLOv8 segmentation (nano version for speed)
 yolo_model = YOLO("yolov8n-seg.pt")  # change to yolov8s-seg.pt for more accuracy
 
-# --- Image Segmentation ---
+# ---------------------------
+# Image Segmentation
+# ---------------------------
 def segment_image(image: Image.Image):
     results = yolo_model.predict(np.array(image))[0]
 
@@ -25,7 +35,9 @@ def segment_image(image: Image.Image):
     overlay_img = Image.fromarray(overlay)
     return (overlay_img, annotations)
 
-# --- Video Segmentation ---
+# ---------------------------
+# Video Segmentation
+# ---------------------------
 def segment_video(video):
     cap = cv2.VideoCapture(video)
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
@@ -54,18 +66,33 @@ def segment_video(video):
     out.release()
     return out_path
 
-# --- Gradio UI ---
+# ---------------------------
+# Text Generation
+# ---------------------------
+def generate_text(prompt):
+    result = text_gen(prompt, max_length=100, num_return_sequences=1)
+    return result[0]["generated_text"]
+
+# ---------------------------
+# Gradio UI
+# ---------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🖼️ YOLOv8 Segmentation (Image + Video)")
+    gr.Markdown("# 🔥 Multi-Modal Playground\nTry out **Text + Image + Video Segmentation** in one app!")
+
+    with gr.Tab("💬 Text Generation"):
+        inp_text = gr.Textbox(label="Enter your prompt")
+        out_text = gr.Textbox(label="Generated text")
+        btn_text = gr.Button("Generate")
+        btn_text.click(generate_text, inputs=inp_text, outputs=out_text)
 
-    with gr.Tab("Image Segmentation"):
+    with gr.Tab("🖼️ Image Segmentation"):
         inp_img = gr.Image(type="pil", label="Upload Image")
         out_img = gr.Image(type="pil", label="Segmented Image")
         out_ann = gr.JSON(label="Annotations")
         btn_img = gr.Button("Run Segmentation")
         btn_img.click(segment_image, inputs=inp_img, outputs=[out_img, out_ann])
 
-    with gr.Tab("Video Segmentation"):
+    with gr.Tab("🎥 Video Segmentation"):
         inp_vid = gr.Video(label="Upload Video")
         out_vid = gr.Video(label="Segmented Video")
         btn_vid = gr.Button("Run Segmentation")