Update app.py

app.py

@@ -18,6 +18,21 @@ def load_model():
 
 model, processor = load_model()
 
+SYSTEM_PROMPT = """You are an AI assistant specialized in analyzing images and videos of code editors like VSCode. Your primary task is to identify and extract code snippets visible in these visual inputs. Follow these guidelines:
+
+1. Focus on recognizing and extracting code, ignoring non-code elements like toolbars, file explorers, or terminals.
+2. If multiple code snippets or files are visible, extract and separate them clearly.
+3. Preserve the syntax, indentation, and structure of the code as seen in the image/video.
+4. If you detect syntax errors or warning highlights in the code, mention them but do not attempt to correct the code.
+5. If the code is partially visible or cut off, extract what you can see and indicate where the code might be incomplete.
+6. Identify the programming language if possible, based on syntax or file extensions visible.
+7. If asked about specific elements (e.g., "What's on line 20?"), focus on that particular part of the code.
+8. Do not invent or assume code that isn't visible in the image/video.
+9. If the image/video doesn't contain code or an IDE, politely inform the user.
+
+Always strive for accuracy in code extraction, as the user may need to directly use or analyze the extracted code."""
+
+
 @spaces.GPU
 def process_image(image, prompt):
     messages = [
@@ -115,25 +130,43 @@ def process_video(video, prompt, max_frames=16, frame_interval=30, max_resolutio
 @spaces.GPU
 def process_content(content, prompt):
     if content is None:
-        return "Please upload an image or video file."
+        return "Please upload an image or video file of a code editor."
+
+    full_prompt = f"{SYSTEM_PROMPT}\n\nUser request: {prompt}"
 
     if content.name.lower().endswith(('.png', '.jpg', '.jpeg')):
-        return process_image(Image.open(content.name), prompt)
+        return process_image(Image.open(content.name), full_prompt)
     elif content.name.lower().endswith(('.mp4', '.avi', '.mov')):
-        return process_video(content, prompt)
+        return process_video(content, full_prompt)
     else:
-        return "Unsupported file type. Please provide an image or video file."
+        return "Unsupported file type. Please provide an image or video file of a code editor."
+
+# Predefined prompts
+PREDEFINED_PROMPTS = [
+    "Extract all visible code from this IDE screenshot.",
+    "What programming language is used in this code editor image?",
+    "Are there any syntax errors highlighted in this VSCode recording?",
+    "Extract the function definition starting at line 15 in this image.",
+    "What are the variable names defined in the visible code?",
+    "Extract and separate all different code snippets or files visible in this IDE recording.",
+    "Is there any commented code in this screenshot? If so, extract it.",
+    "What coding style or convention is being used in this code (e.g., camelCase, snake_case)?",
+    "Are there any import statements or library inclusions visible in this IDE image?",
+    "Extract only the CSS code visible in this multi-file editor screenshot."
+]
 
 # Gradio interface
 iface = gr.Interface(
     fn=process_content,
     inputs=[
-        gr.File(label="Upload Image or Video"),
-        gr.Textbox(label="Enter your prompt")
+        gr.File(label="Upload Image or Video of Code Editor"),
+        gr.Dropdown(choices=PREDEFINED_PROMPTS, label="Select a predefined prompt", type="value"),
+        gr.Textbox(label="Or enter a custom prompt")
     ],
     outputs="text",
-    title="Image and Video Description",
-    description="Upload an image or video and enter a prompt to get a description or analysis.",
+    title="Code Extraction from IDE Screenshots/Recordings",
+    description="Upload an image or video of a code editor and select a predefined prompt or enter a custom one to extract and analyze the code.",
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":