Update app.py

app.py

@@ -28,6 +28,7 @@ def save_captions_to_file(course: str, captions: List[Dict]) -> None:
         file_path = get_caption_file_path(course)
         with open(file_path, 'w', encoding='utf-8') as f:
             json.dump(captions, f, indent=2, ensure_ascii=False)
+        print(f"✓ Saved {len(captions)} captions for {course}")
     except Exception as e:
         print(f"Error saving captions for {course}: {e}")
 
@@ -37,7 +38,9 @@ def load_captions_from_file(course: str) -> List[Dict]:
         file_path = get_caption_file_path(course)
         if file_path.exists():
             with open(file_path, 'r', encoding='utf-8') as f:
-                return json.load(f)
+                captions = json.load(f)
+                print(f"✓ Loaded {len(captions)} existing captions for {course}")
+                return captions
     except Exception as e:
         print(f"Error loading captions for {course}: {e}")
     return []
@@ -131,6 +134,7 @@ app = FastAPI(
 # Global state
 processed_images: Dict[str, Set[str]] = {}  # {course: set(image_names)}
 course_captions: Dict[str, List[Dict]] = {}  # {course: [{image, caption, metadata}]}
+failed_images: Dict[str, Set[str]] = {}  # {course: set(image_names)}
 servers = []
 is_processing = False
 current_processing_task = None
@@ -223,13 +227,15 @@ async def get_processing_status():
     for course in processed_images:
         total = len(processed_images[course])
         processed = len(course_captions.get(course, []))
-        status_info[course] = ProcessingStatus(
-            course=course,
-            total_images=total,
-            processed_images=processed,
-            progress_percent=(processed / total * 100) if total > 0 else 0,
-            status="processing" if processed < total else "completed"
-        )
+        failed = len(failed_images.get(course, set()))
+        status_info[course] = {
+            "course": course,
+            "total_images": total,
+            "processed_images": processed,
+            "failed_images": failed,
+            "progress_percent": (processed / total * 100) if total > 0 else 0,
+            "status": "completed" if processed + failed >= total else "processing"
+        }
     return status_info
 
 @app.post("/processing/start")
@@ -291,6 +297,8 @@ async def delete_captions(course: str):
                 del processed_images[course]
             if course in course_captions:
                 del course_captions[course]
+            if course in failed_images:
+                del failed_images[course]
             return {"message": f"Captions for {course} deleted"}
         else:
             raise HTTPException(status_code=404, detail=f"No captions found for {course}")
@@ -417,21 +425,24 @@ async def submit_to_dataset(course: str, metadata_list: List[Dict]):
                 json=payload
             ) as resp:
                 result = await resp.json()
-                print(f"Dataset submission result for {course}: {result}")
+                print(f"✓ Dataset submission result for {course}: {result}")
                 return result
         except Exception as e:
-            print(f"Error submitting to dataset: {e}")
+            print(f"✗ Error submitting to dataset: {e}")
             return None
 
 async def process_course(course: str, servers: List[CaptionServer]):
-    """Process all images in a course using available servers with retry logic"""
+    """Process all images in a course using available servers with proper retry logic"""
+    # Initialize course tracking
     if course not in processed_images:
         processed_images[course] = set()
     if course not in course_captions:
-        # Load any existing captions from file
         course_captions[course] = load_captions_from_file(course)
         # Update processed images set from loaded captions
-        processed_images[course].update(cap['image'] for cap in course_captions[course])
+        for cap in course_captions[course]:
+            processed_images[course].add(cap['image'])
+    if course not in failed_images:
+        failed_images[course] = set()
     
     # Get list of images
     images = await fetch_course_images(course)
@@ -442,18 +453,21 @@ async def process_course(course: str, servers: List[CaptionServer]):
     print(f"\nProcessing {len(images)} images for course {course}")
     
     # Track images that need processing with retry count (5 retries)
-    pending_images = {
-        img['filename']: {'image': img, 'retries': 0, 'max_retries': 5}
-        for img in images 
-        if img['filename'] not in processed_images[course]
-    }
+    pending_images = {}
+    for img in images:
+        filename = img['filename']
+        if filename not in processed_images[course] and filename not in failed_images[course]:
+            pending_images[filename] = {'image': img, 'retries': 0, 'max_retries': 5}
     
     if not pending_images:
-        print(f"All images already processed for course {course}")
+        print(f"All images already processed or failed for course {course}")
+        print(f"- Processed: {len(processed_images[course])}, Failed: {len(failed_images[course])}")
         return
     
-    failed_images = set()
-    total_images = len(images)
+    print(f"Images to process: {len(pending_images)} (already processed: {len(processed_images[course])}, failed: {len(failed_images[course])})")
+    
+    batch_size = len([s for s in servers if not s.busy])
+    processed_in_this_run = 0
     
     while pending_images and is_processing:
         # Create tasks for each available server
@@ -468,7 +482,7 @@ async def process_course(course: str, servers: List[CaptionServer]):
                 
                 # Assign this image to the server
                 tasks.append(process_image(server, course, img))
-                assigned_images.append((filename, img))
+                assigned_images.append((filename, img, img_data['retries']))
                 # Remove from pending temporarily while it's being processed
                 del pending_images[filename]
                 
@@ -482,52 +496,63 @@ async def process_course(course: str, servers: List[CaptionServer]):
         
         # Handle results and retry logic
         has_new_results = False
-        for (filename, img), result in zip(assigned_images, results):
+        for (filename, img, current_retries), result in zip(assigned_images, results):
             if result:
                 # Success - image was processed
                 processed_images[course].add(filename)
                 course_captions[course].append(result)
                 has_new_results = True
+                processed_in_this_run += 1
                 print(f"✓ Successfully processed {filename}")
             else:
                 # Failure - check if we should retry
-                current_retries = assigned_images[0][0] if assigned_images else 0
-                img_data = {'image': img, 'retries': current_retries + 1, 'max_retries': 5}
-                
-                if img_data['retries'] <= img_data['max_retries']:
-                    # Put back in pending for retry
-                    pending_images[filename] = img_data
-                    print(f"↻ Retry {img_data['retries']}/{img_data['max_retries']} for {filename}")
+                if current_retries < 5:  # max_retries
+                    # Put back in pending for retry with incremented retry count
+                    pending_images[filename] = {
+                        'image': img, 
+                        'retries': current_retries + 1, 
+                        'max_retries': 5
+                    }
+                    print(f"↻ Retry {current_retries + 1}/5 for {filename}")
                 else:
                     # Max retries exceeded, mark as failed
-                    failed_images.add(filename)
-                    print(f"✗ Failed to process {filename} after {img_data['max_retries']} retries")
+                    failed_images[course].add(filename)
+                    print(f"✗ Failed to process {filename} after 5 retries")
         
         # Save progress after each batch with new results
         if has_new_results:
             save_captions_to_file(course, course_captions[course])
         
         # Show progress
+        total = len(images)
         done = len(processed_images[course])
+        failed_count = len(failed_images[course])
         pending_count = len(pending_images)
-        failed_count = len(failed_images)
-        progress_percent = (done / total_images * 100) if total_images > 0 else 0
-        print(f"\rProgress: {done}/{total_images} images ({progress_percent:.1f}%) - {pending_count} pending, {failed_count} failed", end="")
+        progress_percent = (done / total * 100) if total > 0 else 0
+        
+        print(f"\rProgress: {done}/{total} ({progress_percent:.1f}%) - {pending_count} pending, {failed_count} failed, {processed_in_this_run} new", end="", flush=True)
         
-        # Check if we're done (all images either processed or permanently failed)
-        if not pending_images:
-            if failed_images:
-                print(f"\nCourse {course} completed with {len(failed_images)} failed images: {list(failed_images)}")
-            else:
-                print(f"\nCourse {course} complete, submitting to dataset...")
-                await submit_to_dataset(course, course_captions[course])
-            break
-
         # Small delay to prevent overwhelming the servers
         await asyncio.sleep(0.5)
+    
+    # Final status for this course
+    total = len(images)
+    done = len(processed_images[course])
+    failed_count = len(failed_images[course])
+    
+    if done + failed_count >= total:
+        if failed_count > 0:
+            print(f"\n✓ Course {course} completed with {failed_count} failed images")
+        else:
+            print(f"\n✓ Course {course} fully completed")
+            # Submit to dataset only when fully completed
+            print(f"Submitting {len(course_captions[course])} captions to dataset...")
+            await submit_to_dataset(course, course_captions[course])
+    else:
+        print(f"\n→ Course {course} partially completed: {done}/{total} processed, {failed_count} failed")
 
 async def processing_loop(specific_courses: Optional[List[str]] = None, continuous: bool = True):
-    """Main processing loop - same as original main() function"""
+    """Main processing loop with proper error handling"""
     global is_processing
     
     # Get model information and verify Florence-2-large availability
@@ -566,14 +591,22 @@ async def processing_loop(specific_courses: Optional[List[str]] = None, continuo
         print()
     
     start_time = time.time()
+    iteration = 0
     
     while is_processing:
         try:
+            iteration += 1
+            print(f"\n{'='*50}")
+            print(f"Processing Iteration {iteration}")
+            print(f"{'='*50}")
+            
             # Get available courses
             if specific_courses:
                 courses = specific_courses
+                print(f"Processing specific courses: {courses}")
             else:
                 courses = await fetch_courses()
+                print(f"Found {len(courses)} courses")
                 
             if not courses:
                 print("No courses found, waiting...")
@@ -582,24 +615,24 @@ async def processing_loop(specific_courses: Optional[List[str]] = None, continuo
                 await asyncio.sleep(10)
                 continue
             
-            print(f"Found {len(courses)} courses")
-            
             # Process each course with all available servers
             for course in courses:
                 if not is_processing:
                     break
+                    
+                print(f"\n--- Processing course: {course} ---")
                 await process_course(course, processing_servers)
                 
-                # Show server stats
-                print("\nServer Stats:")
-                total_processed = sum(s.total_processed for s in processing_servers)
-                elapsed = time.time() - start_time
-                if elapsed > 0:
-                    print(f"Total images processed: {total_processed}")
-                    print(f"Overall speed: {total_processed/elapsed:.2f} fps")
-                    for s in processing_servers:
-                        print(f"- {s.url}: {s.total_processed} images, {s.fps:.2f} fps")
-                print()
+            # Show server stats
+            print("\nServer Stats:")
+            total_processed = sum(s.total_processed for s in processing_servers)
+            elapsed = time.time() - start_time
+            if elapsed > 0:
+                print(f"Total images processed: {total_processed}")
+                print(f"Overall speed: {total_processed/elapsed:.2f} fps")
+                for s in processing_servers:
+                    print(f"- {s.url}: {s.total_processed} images, {s.fps:.2f} fps")
+            print()
             
             if not continuous:
                 print("One-time processing completed")
@@ -613,7 +646,9 @@ async def processing_loop(specific_courses: Optional[List[str]] = None, continuo
             print("Processing cancelled")
             break
         except Exception as e:
-            print(f"Error in processing loop: {e}")
+            print(f"Error in processing loop: {str(e)}")
+            import traceback
+            traceback.print_exc()
             await asyncio.sleep(10)
     
     is_processing = False