Update app.py

app.py

@@ -11,52 +11,52 @@ DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 TORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float16
 
 
-def get_step_info(step_name):
+def get_step_info(step_number):
     """Returns detailed information about a manufacturing step."""
     step_details = {
-        "Step 1": {
+        1: {
             "Name": "Bead Insertion",
             "Standard Time": "4 seconds",
             "Analysis": "Observe the bead placement process. If the insertion exceeds 4 seconds, identify potential issues such as missing beads, technician errors, or machinery malfunction."
         },
-        "Step 2": {
+        2: {
             "Name": "Inner Liner Apply",
             "Standard Time": "4 seconds",
             "Analysis": "Check for manual intervention during the inner layer application. If adjustment is required, it may indicate improper alignment or issues with the layer material."
         },
-        "Step 3": {
+        3: {
             "Name": "Ply1 Apply",
             "Standard Time": "4 seconds",
             "Analysis": "Determine if the technician is manually adjusting the first ply. Manual intervention might suggest improper ply placement or machine misalignment."
         },
-        "Step 4": {
+        4: {
             "Name": "Bead Set",
             "Standard Time": "8 seconds",
             "Analysis": "Observe the bead setting process. Delays may result from bead misalignment, machine pauses, or lack of technician involvement."
         },
-        "Step 5": {
+        5: {
             "Name": "Turnup",
             "Standard Time": "4 seconds",
             "Analysis": "Examine the turnup step for any technician involvement or pauses in machine operation. Reasons for delays might include material misalignment or equipment issues."
         },
-        "Step 6": {
+        6: {
             "Name": "Sidewall Apply",
             "Standard Time": "14 seconds",
             "Analysis": "If a technician is repairing the sidewall, this may indicate material damage or improper initial application. Look for signs of excessive manual handling."
         },
-        "Step 7": {
+        7: {
             "Name": "Sidewall Stitching",
             "Standard Time": "5 seconds",
             "Analysis": "Observe the stitching process. Delays could occur due to machine speed inconsistencies or technician intervention for correction."
         },
-        "Step 8": {
+        8: {
             "Name": "Carcass Unload",
             "Standard Time": "7 seconds",
             "Analysis": "Ensure a technician is present for the carcass unload. If absent, note their return time and identify potential reasons for their absence."
         }
     }
     
-    return step_details.get(step_name, {"Error": "Invalid step name. Please provide a valid step number."})
+    return step_details.get(step_number, {"Error": "Invalid step number. Please provide a valid step number."})
 
 
 def load_video(video_data, strategy='chat'):
@@ -140,11 +140,11 @@ def predict(prompt, video_data, temperature, model, tokenizer):
     
     return response
 
-def get_analysis_prompt(step_name):
+def get_analysis_prompt(step_number):
     """Constructs the prompt for analyzing delay reasons based on the selected step."""
     return f"""You are an AI expert system specialized in analyzing manufacturing processes and identifying production delays in tire manufacturing. Your role is to accurately classify delay reasons based on visual evidence from production line footage.
 Task Context:
-You are analyzing video footage from Step {step_name} of a tire manufacturing process where a delay has been detected. 
+You are analyzing video footage from Step {step_number} of a tire manufacturing process where a delay has been detected. 
 Required Analysis:
  Carefully observe the video for visual cues indicating production interruption.
  If no person is visible in any of the frames, the reason probably might be due to his absence.
@@ -160,7 +160,7 @@ Important: Base your analysis solely on visual evidence from the video. Focus on
 
 model, tokenizer = load_model()
 
-def inference(video, step_name):
+def inference(video, step_number):
     """Analyzes video to predict possible issues based on the manufacturing step."""
     try:
         if not video:
@@ -197,8 +197,8 @@ def create_interface():
         
         gr.Examples(
             examples=[
-                ["7838_step2_2_eval.mp4", "Step 2"],
-                ["7838_step6_2_eval.mp4", "Step 6"]
+                ["7838_step2_2_eval.mp4", 2],
+                ["7838_step6_2_eval.mp4", 6]
             ],
             inputs=[video, step_number],
             cache_examples=False
@@ -214,4 +214,4 @@ def create_interface():
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.queue().launch(share=True)
+    demo.queue().launch(share=True)