import numpy as np
import cv2
import torch
import torch.nn as nn

class ScreenClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1, bias=False), nn.BatchNorm2d(16), nn.ReLU(True), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.ReLU(True), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True), nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Sequential(nn.Flatten(), nn.Dropout(0.3), nn.Linear(64, 1))

    def forward(self, x):
        return self.classifier(self.features(x))

# Load
model = ScreenClassifier()
model.load_state_dict(torch.load("screen_classifier_best.pth", map_location="cpu", weights_only=True))
model.eval()

# Predict from OpenCV frame
frame = cv2.imread("phone_screen.jpg")
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
resized = cv2.resize(gray, (64, 64), interpolation=cv2.INTER_AREA)
tensor = torch.from_numpy(resized.astype(np.float32)).div(255.0)
tensor = (tensor - 0.5) / 0.5
tensor = tensor.unsqueeze(0).unsqueeze(0)

with torch.no_grad():
    prob = torch.sigmoid(model(tensor).squeeze()).item()

label = "ON" if prob >= 0.5 else "OFF"
confidence = prob if label == "ON" else 1.0 - prob
print(f"{label} (confidence: {confidence:.1%})")

Training

Trained on synthetic data (23473 params) with domain randomization.

ON screens: 10 UI layout styles (light/dark apps, chat, media, settings, browser, etc.)
OFF screens: 9 dark glass variations (glare, reflections, fingerprints, ambient lighting)

Files

screen_classifier_best.pth — PyTorch state dict
screen_classifier_best.pt — TorchScript (deploy without class definition)
metrics.json — Training metrics
screen_classifier.py — Full training + inference code

dhruvguptaa/screen-on-off-classifier

Screen ON/OFF Classifier

Performance

Usage

Training

Files

dhruvguptaa/
screen-on-off-classifier