001/*
002 * Copyright (C) Photon Vision.
003 *
004 * This program is free software: you can redistribute it and/or modify
005 * it under the terms of the GNU General Public License as published by
006 * the Free Software Foundation, either version 3 of the License, or
007 * (at your option) any later version.
008 *
009 * This program is distributed in the hope that it will be useful,
010 * but WITHOUT ANY WARRANTY; without even the implied warranty of
011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
012 * GNU General Public License for more details.
013 *
014 * You should have received a copy of the GNU General Public License
015 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
016 */
017
018package org.photonvision.vision.objects;
019
020import java.util.ArrayList;
021import java.util.List;
022import org.opencv.core.Core;
023import org.opencv.core.Mat;
024import org.opencv.core.Rect2d;
025import org.opencv.core.Scalar;
026import org.opencv.core.Size;
027import org.opencv.imgproc.Imgproc;
028import org.photonvision.vision.pipe.impl.NeuralNetworkPipeResult;
029
030public class Letterbox {
031    double dx;
032    double dy;
033    double scale;
034
035    public Letterbox(double dx, double dy, double scale) {
036        this.dx = dx;
037        this.dy = dy;
038        this.scale = scale;
039    }
040
041    /**
042     * Resize the frame to the new shape and "letterbox" it.
043     *
044     * <p>Letterboxing is the process of resizing an image to a new shape while maintaining the aspect
045     * ratio of the original image. The new image is padded with a color to fill the remaining space.
046     *
047     * @param frame
048     * @param letterboxed
049     * @param newShape
050     * @param color
051     * @return
052     */
053    public static Letterbox letterbox(Mat frame, Mat letterboxed, Size newShape, Scalar color) {
054        // from https://github.com/ultralytics/yolov5/issues/8427#issuecomment-1172469631
055        var frameSize = frame.size();
056        var r = Math.min(newShape.height / frameSize.height, newShape.width / frameSize.width);
057
058        var newUnpad = new Size(Math.round(frameSize.width * r), Math.round(frameSize.height * r));
059
060        if (!(frameSize.equals(newUnpad))) {
061            Imgproc.resize(frame, letterboxed, newUnpad, Imgproc.INTER_LINEAR);
062        } else {
063            frame.copyTo(letterboxed);
064        }
065
066        var dw = newShape.width - newUnpad.width;
067        var dh = newShape.height - newUnpad.height;
068
069        dw /= 2;
070        dh /= 2;
071
072        int top = (int) (Math.round(dh - 0.1f));
073        int bottom = (int) (Math.round(dh + 0.1f));
074        int left = (int) (Math.round(dw - 0.1f));
075        int right = (int) (Math.round(dw + 0.1f));
076        Core.copyMakeBorder(
077                letterboxed, letterboxed, top, bottom, left, right, Core.BORDER_CONSTANT, color);
078
079        return new Letterbox(dw, dh, r);
080    }
081
082    /**
083     * Resizes the detections to the original frame size.
084     *
085     * @param unscaled The detections to resize
086     * @return The resized detections
087     */
088    public List<NeuralNetworkPipeResult> resizeDetections(List<NeuralNetworkPipeResult> unscaled) {
089        var ret = new ArrayList<NeuralNetworkPipeResult>();
090
091        for (var t : unscaled) {
092            var scale = 1.0 / this.scale;
093            var boundingBox = t.bbox;
094            double x = (boundingBox.x - this.dx) * scale;
095            double y = (boundingBox.y - this.dy) * scale;
096            double width = boundingBox.width * scale;
097            double height = boundingBox.height * scale;
098
099            ret.add(
100                    new NeuralNetworkPipeResult(new Rect2d(x, y, width, height), t.classIdx, t.confidence));
101        }
102
103        return ret;
104    }
105}