(feat): add ResNet18Vision (1D); register; inference --arch supports it

models/registry.py

@@ -2,13 +2,13 @@
 from typing import Callable, Dict
 from models.figure2_cnn import Figure2CNN
 from models.resnet_cnn import ResNet1D
-# from models.resnet18_vision import ResNet18Vision # (Step 2)
+from models.resnet18_vision import ResNet18Vision 
 
 # Internal registry of model builders keyed by short name.
 _REGISTRY: Dict[str, Callable[[int], object]] = {
     "figure2": lambda L: Figure2CNN(input_length=L),
     "resnet": lambda L: ResNet1D(input_length=L),
-    # "resnet18vision": lambda L: ResNet18Vision(input_length=L)
+    "resnet18vision": lambda L: ResNet18Vision(input_length=L)
 }
 
 def choices():

models/resnet18_vision.py

@@ -0,0 +1,78 @@
+# models/resnet18_vision.py
+# 1D ResNet-18 style model for spectra: input (B, 1, L)
+import torch
+import torch.nn as nn
+from typing import Callable, List
+
+class BasicBlock1D(nn.Module):
+    expansion = 1
+    def __init__(self, in_planes: int, planes: int, stride: int = 1, downsample: nn.Module | None = None):
+        super().__init__()
+        self.conv1 = nn.Conv1d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1   = nn.BatchNorm1d(planes)
+        self.relu  = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv1d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2   = nn.BatchNorm1d(planes)
+        self.downsample = downsample
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        out = self.relu(out)
+        return out
+
+def _make_layer(block: Callable[..., nn.Module], in_planes: int, planes: int, blocks: int, stride: int) -> nn.Sequential:
+    downsample = None
+    if stride != 1 or in_planes != planes * block.expansion:
+        downsample = nn.Sequential(
+            nn.Conv1d(in_planes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
+            nn.BatchNorm1d(planes * block.expansion),
+        )
+    layers: List[nn.Module] = [block(in_planes, planes, stride, downsample)]
+    in_planes = planes * block.expansion
+    for _ in range(1, blocks):
+        layers.append(block(in_planes, planes))
+    return nn.Sequential(*layers)
+
+class ResNet18Vision(nn.Module):
+    def __init__(self, input_length: int = 500, num_classes: int = 2):
+        super().__init__()
+        # 1D stem
+        self.conv1 = nn.Conv1d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1   = nn.BatchNorm1d(64)
+        self.relu  = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+
+        # ResNet-18: 2 blocks per layer
+        self.layer1 = _make_layer(BasicBlock1D, 64,  64, blocks=2, stride=1)
+        self.layer2 = _make_layer(BasicBlock1D, 64, 128, blocks=2, stride=2)
+        self.layer3 = _make_layer(BasicBlock1D, 128, 256, blocks=2, stride=2)
+        self.layer4 = _make_layer(BasicBlock1D, 256, 512, blocks=2, stride=2)
+
+        # Global pooling + classifier
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.fc = nn.Linear(512 * BasicBlock1D.expansion, num_classes)
+
+        # Kaiming init
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm1d, nn.GroupNorm)):
+                nn.init.ones_(m.weight); nn.init.zeros_(m.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, 1, L)
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.avgpool(x)     # (B, C, 1)
+        x = torch.flatten(x, 1) # (B, C)
+        x = self.fc(x)          # (B, num_classes)
+        return x

scripts/run_inference.py

@@ -9,9 +9,11 @@ import logging
 
 import numpy as np
 import torch
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-from models.figure2_cnn import Figure2CNN
 from scripts.preprocess_dataset import resample_spectrum, label_file
+from models.registry import choices as model_choices, build as build_model
+
 
 
 # =============================================
@@ -49,6 +51,8 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Run inference on a single Raman spectrum (.txt file)."
     )
+    parser.add_argument("--arch", type=str, default="figure2", choices=model_choices(),
+                    help="Model architecture (must match the provided weights).")  # NEW
     parser.add_argument(
         "--target-len", type=int, required=True,
         help="Target length to match model input"
@@ -96,18 +100,17 @@ if __name__ == "__main__":
 
         data = resample_spectrum(x_raw, y_raw, target_len=args.target_len)
         # Shape = (1, 1, target_len) — valid input for Raman inference
-        input_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+        input_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(DEVICE)
+
 
-        # 2. Load Model
-        model = Figure2CNN(
-            input_length=args.target_len,
-            input_channels=1
-        )
+        # 2. Load Model (via shared model registry)
+        model = build_model(args.arch, args.target_len).to(DEVICE)
         if args.model != "random":
-            model.load_state_dict(
-                torch.load(args.model, map_location="cpu", weights_only=True)
-            )
+            state = torch.load(args.model, map_location="cpu") # broad compatibility
+            model.load_state_dict(state)
         model.eval()
+        
+        
 
         # 3. Inference
         with torch.no_grad():