001    package calhoun.analysis.crf.executables.viewer;
002    
003    import java.awt.Color;
004    import java.awt.Component;
005    
006    import javax.swing.JTable;
007    import javax.swing.table.AbstractTableModel;
008    import javax.swing.table.DefaultTableCellRenderer;
009    
010    import calhoun.analysis.crf.Conrad;
011    import calhoun.analysis.crf.CRFInference.InferenceResult;
012    import calhoun.analysis.crf.io.TrainingSequence;
013    
014    public class ViterbiTableModel extends AbstractTableModel {
015            private static final long serialVersionUID = 4315252264878822897L;
016            Conrad crfModel;
017            TrainingSequence seq;
018            int[] viterbiPath;
019            double[] bestScore;
020            int nStates;
021    
022            public ViterbiTableModel(Conrad crfModel, TrainingSequence seq) {
023                    this.crfModel = crfModel;
024                    this.seq = seq;
025                    nStates = crfModel.getModel().getNumStates();
026                    InferenceResult result = crfModel.predict(seq);
027                    viterbiPath = result.hiddenStates;
028                    bestScore = null;
029            }
030    
031            @Override
032            public String getColumnName(int col) {
033                    return (col == 0) ? "Position" : "    " + Integer.toString(col) + "    ";
034            }
035    
036            public int getRowCount() {
037                    return crfModel.getModel().getNumStates() + 2;
038            }
039    
040            public int getColumnCount() {
041                    return seq.length() + 1;
042            }
043    
044            @Override
045            public Class getColumnClass(int col) {
046                    return String.class;
047            }
048    
049            public Object getValueAt(int row, int col) {
050                    switch (row) {
051                    case 0:
052                            return col == 0 ? "Input" : seq.getInputSequence().getX(col - 1).toString();
053                    case 1:
054                            return col == 0 ? "Label" : Integer.toString(seq.getY(col - 1));
055                    default:
056                            return col == 0 ? crfModel.getModel().getStateName(row - 2) : Double.toString(bestScore[(row - 2) * nStates + col - 1]);
057                    }
058            }
059    
060            @Override
061            public boolean isCellEditable(int row, int col) {
062                    return false;
063            }
064    
065            @Override
066            public void setValueAt(Object value, int row, int col) {
067                    // this 1 line was screwing up the sort for almost a year.
068                    //m_tableSorter.fireTableDataChanged();
069            }
070    
071            public class ViterbiCellRenderer extends DefaultTableCellRenderer {
072                    private static final long serialVersionUID = -1156070614840797662L;
073    
074                    @Override
075                    public Component getTableCellRendererComponent(JTable table, Object value, boolean isSelected, boolean hasFocus, int row, int column) {
076                            super.getTableCellRendererComponent(table, value , isSelected, hasFocus, row, column);
077                            if(row > 1 && column > 0 && (row-2) == viterbiPath[column-1]) {
078                                    setBackground(Color.RED);
079                            }
080                            else {
081                                    setBackground(Color.WHITE);
082                            }
083                            return this;
084                    }
085            }
086    }