Allow class_names.txt upload

Author: probicheaux

roboflow/util/model_processor.py

@@ -250,6 +250,10 @@ def _process_rfdetr(model_type: str, model_path: str, filename: str) -> str:
     if pt_file is None:
         raise RuntimeError("No .pt model file found in the provided path")
 
+
+    class_names = get_classnames_txt_for_rfdetr(model_path, pt_file)
+
+        
     # Copy the .pt file to weights.pt if not already named weights.pt
     if pt_file != "weights.pt":
         shutil.copy(os.path.join(model_path, pt_file), os.path.join(model_path, "weights.pt"))
@@ -269,6 +273,26 @@ def _process_rfdetr(model_type: str, model_path: str, filename: str) -> str:
 
     return zip_file_name
 
+def get_classnames_txt_for_rfdetr(model_path: str, pt_file: str) -> list[str]:
+    class_names_path = os.path.join(model_path, "class_names.txt")
+    if os.path.exists(class_names_path):
+        return class_names_path
+    
+    model = torch.load(os.path.join(model_path, pt_file), map_location="cpu", weights_only=False)
+    args = vars(model["args"])
+    if "class_names" in args:
+        with open(class_names_path, "w") as f:
+            for class_name in args["class_names"]:
+                f.write(class_name + "\n")
+        return class_names_path
+
+    raise FileNotFoundError(
+        f"No class_names.txt file found in model path {model_path}.\n"
+        f"This should only happen on rfdetr models trained before version 1.1.0.\n"
+        f"Please re-train your model with the latest version of the rfdetr library, or\n"
+        f"please create a class_names.txt file in the model path with the class names in new lines in the order of the classes in the model.\n"
+    )
+
 
 def _process_huggingface(
     model_type: str, model_path: str, filename: str = "fine-tuned-paligemma-3b-pt-224.f16.npz"