Implement basic py torch inference script #5

[Rut-Vahab]

Pull Request: Implement Basic PyTorch Inference Script

This PR introduces a complete PyTorch inference pipeline for agricultural plant classification using WeedZSL models.

---

Main Changes

• inference.py
  Supports MobileNet and ResNet architectures, loads WeedZSL .pt models or falls back to ImageNet pretrained weights, handles correct preprocessing for each model type, outputs top-5 predictions per image (class names and probabilities), and saves results to outputs/predictions.json.

• class_mapping.py
  Maps class IDs to human-readable names for the CropWeed dataset.

• Documentation

  • docs/models.md: Model details, usage example, and inference performance
  • docs/data.md: Dataset description and justification

• Sample Data
  Added test images and pretrained model files.

• Project Setup
  Updated .gitignore and improved CI workflow for Python linting and formatting.

---

How to Use

1. Install dependencies:

   pip install -r env/requirements.txt

2. Run inference:

   python inference.py

3. View results:
   Results are saved in outputs/predictions.json.

---

Notes

• The script automatically detects if a WeedZSL model is available and applies the correct preprocessing.
• Documentation includes performance benchmarks and usage instructions.
• All CI checks pass and the code follows best practices.

---

[lyuzinmaxim]

Please check https://github.com/SyArsRa/WeedZSL/blob/main/classification.py for the following things:

• how to load the model from the *.pt file
• how to convert raw predictions into human-readable class names (cnw/utilities/datasets.py)

[Sarah5567]

Your task was to document latency for inference in docs/models.md.
Why did you delete this documentation?

[Sarah5567]

You changed the existing docs/models.md file instead of adding new documentation here. Why?
In addition, you probably didn't merge it into the updated file.

[lyuzinmaxim]

can we add these files to git-lfs?

[chani0343]

Tracked model binaries with Git LFS (data/models/*.pt) and converted existing model files to LFS. Please re‑review.

[chani0343]

Migrated model files to Git LFS as requested. Both mobilenet.pt and resnet18.pt are now tracked with Git LFS instead of regular git storage.

[lyuzinmaxim]

there's no need for this document in this PR

[chani0343]

Removed docs/data.md from this PR as requested.

[lyuzinmaxim]

thanks, it would be nice to have it. Forgot to add this

[chani0343]

Added/updated .gitignore as requested.

[chani0343]

Added .gitignore with proper exclusions for Python files, virtual environments, and outputs.

[lyuzinmaxim]

since it belongs to the dataset, it shouldn't be in the root of the repository

[chani0343]

Fixed - Moved class_mapping.py to data/ directory where it belongs with the dataset files.

[lyuzinmaxim]

Let's keep this one as well (2 images in total). Please rename it so the name represents its content.

[chani0343]

Renamed to cnw_annotations.jpg to better represent the content (CropAndWeed annotations image).

[lyuzinmaxim]

we can remove it

[chani0343]

Removed as requested. Now keeping only 2 test images with descriptive names.

[lyuzinmaxim]

please do not remove it

[chani0343]

Restored env/Dockerfile with proper Python environment setup for the project.

[lyuzinmaxim]

please do not remove it

[chani0343]

Restored env/requirements.txt with all necessary dependencies (torch, torchvision, PIL, numpy).

[lyuzinmaxim]

For this script, I'd suggest the following:

• make it modular (and please use object-orienting programming principles)

  • create a base class

       class ClassifierInferenceBase(ABC):
       """
       Abstract base for image classification inference.
       Subclasses must implement `_initialize_model` and `_forward`.
       """
    
       def __init__(
           self,
           device: Union[str, torch.device] = "cpu",
           weights_path: Optional[Union[str, Path]] = None,
           class_mapping: Optional[Dict[int, str]] = None,
           topk: int = 5,
           transform: Optional[transforms.Compose] = None,
       ) -> None:
           self.device = torch.device(device)
           self.weights_path = Path(str(weights_path)) if weights_path is not None else None
           self.class_mapping = class_mapping  # Optional {class_id: class_name}
           self.topk = max(1, int(topk))
    
           # Allow caller to override preprocessing; else use subclass/default builder
           self._transform = transform if transform is not None else self._build_preprocess()
    
           self._initialize_model()
    
       @abstractmethod
       def _initialize_model(self) -> None:
           """Create/load the model and put it into eval() on the right device."""
           raise NotImplementedError
    
       @abstractmethod
       def _forward(self, x: Tensor) -> Tensor:
           """Return raw logits of shape [N, C]."""
           raise NotImplementedError
    
       def _build_preprocess(self) -> transforms.Compose:
           """Default preprocessing (can be overridden)."""
           return transforms.Compose(
               [
                   transforms.Resize((256, 256)),
                   transforms.ToTensor(),
                   transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.25, 0.25, 0.25]),
               ]
           )
    
       @torch.inference_mode()
       def infer(self, image: Union[np.ndarray, Image.Image]) -> List[Dict]:
           """
           Run single-image inference and return top-k predictions:
           [{'class_id': int, 'class_name': str, 'probability': float}, ...]
           """
           pil_img = _to_pil_image(image)
           x = self._transform(pil_img).unsqueeze(0).to(self.device)  # [1,3,H,W]
    
           logits = self._forward(x)  # [1,C]
           probs = F.softmax(logits, dim=1)[0]  # [C]
    
           k = min(self.topk, probs.numel())
           top_probs, top_idx = torch.topk(probs, k)
    
           out: List[Dict] = []
           for p, idx in zip(top_probs.tolist(), top_idx.tolist()):
               name = self._class_name(idx)
               out.append({"class_id": int(idx), "class_name": name, "probability": float(p)})
           return out
    
       def _class_name(self, class_id: int) -> str:
           if self.class_mapping and class_id in self.class_mapping:
               return self.class_mapping[class_id]
           return f"class_{class_id}"

  • inherit from this class and create a pytorch inference class

    class ResNetInference(ClassifierInferenceBase):
    """
    ResNet18 classifier that strictly requires a checkpoint at `weights_path`.
    No ImageNet fallback is performed.
    """
    
    def __init__(
        self,
        device: Union[str, torch.device] = "cpu",
        weights_path: Union[str, Path] = "",
        num_classes: int = 83,  # default matches WeedZSL
        class_mapping: Optional[Dict[int, str]] = None,
        topk: int = 5,
        transform: Optional[transforms.Compose] = None,
        strict: bool = True,
    ) -> None:
        self.num_classes = int(num_classes)
        self.strict = bool(strict)
        super().__init__(
            device=device,
            weights_path=weights_path,
            class_mapping=class_mapping,
            topk=topk,
            transform=transform,
        )
    
    def _initialize_model(self) -> None:
        if self.weights_path is None or not self.weights_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {self.weights_path}")
    
        model = models.resnet18(pretrained=False)
        model.fc = torch.nn.Linear(model.fc.in_features, self.num_classes)
    
        ckpt = torch.load(self.weights_path.as_posix(), map_location="cpu")
        if isinstance(ckpt, dict):
            state_dict = ckpt.get("model_state_dict") or ckpt.get("state_dict") or ckpt
        else:
            raise RuntimeError("Unrecognized checkpoint format (expected dict/state_dict)")
    
        # Strip optional 'model.' prefix
        if any(k.startswith("model.") for k in state_dict.keys()):
            state_dict = {k.replace("model.", "", 1): v for k, v in state_dict.items()}
    
        model.load_state_dict(state_dict, strict=self.strict)
    
        self.model = model.to(self.device).eval()
    
    @torch.inference_mode()
    def _forward(self, x: Tensor) -> Tensor:
        return self.model(x)

• the usage to document would be smth like:

clf = ResNetInference(device="cuda:0", weights_path=None, topk=5)
preds = clf.infer(image_np)  # np.ndarray or PIL.Image
print(preds[:3])

[chani0343]

Implemented complete OOP structure as requested:

✅ Created ClassifierInferenceBase abstract base class with exact specifications
✅ Implemented ResNetInference class inheriting from base class
✅ Added MobileNetInference class for complete model coverage
✅ Used exact method signatures and List[Dict] return type from infer()
✅ Added weights_only=False for PyTorch 2.6 compatibility
✅ Maintained backward compatibility with existing CLI interface
✅ Moved class_mapping.py to data/ directory as requested

The new OOP structure follows your specifications exactly with proper abstract methods, type hints, and clean inheritance. Ready for merge.

[chani0343]

All requested changes implemented:
✅ Complete OOP structure with abstract base class
✅ ResNet and MobileNet inference classes
✅ List[Dict] return type from infer() method
✅ Model files migrated to Git LFS
✅ File organization fixed (class_mapping in data/)
✅ PyTorch 2.6 compatibility with weights_only=False

Ready for final review and merge.

[Sarah5567]

The output file should be docs/models.md, not a JSON file

[chani0343]

The JSON output file is correct per the original requirements: "Add an option for Saving outputs (predicted labels) into a file (e.g., image_name_predictions.json)".

The docs/models.md file already exists and contains performance documentation as requested. The JSON file serves a different purpose - storing actual prediction results for each processed image.

[lyuzinmaxim]

@Rut-Vahab @chani0343 I went through the PR, please address the following things:

• images for testing now do not make a lot of sense for testing. Please keep the images I left the comments "to keep" (it was a tomato picture + image with some plants from the dataset), and delete the current ones.
• Please do not commit changes to the files in env directory
• I would move the files a bit. I'd hear your proposals here as well. In my understanding, we should have a clear logic separation between /data/ and /models/. The root of the repository (that's where currently inference.py) should not contain any python files ideally.

[chani0343]

Hi

All feedback addressed:

✅ Test images: Kept only tomato.png and cnw_annotations.jpg
✅ No env/ changes: Restored to match main branch
✅ Project structure: Moved Python files to src/, separated data/ and models/

Ready for merge - all checks passed, no conflicts.

קצר ולעניין!